import json
import asyncio
from channels.generic.websocket import AsyncWebsocketConsumer
from .models import Host
from .ssh_handler import SSHHandler


class SSHConsumer(AsyncWebsocketConsumer):
    """SSH WebSocket消费者"""
    
    async def connect(self):
        """WebSocket连接"""
        self.host_id = self.scope['url_route']['kwargs'].get('host_id')
        self.ssh_handler = None
        self.read_task = None
        
        await self.accept()
        
        try:
            # 获取主机信息
            host = await self.get_host(self.host_id)
            if not host:
                await self.send(text_data=json.dumps({
                    'type': 'error',
                    'message': '主机不存在'
                }))
                await self.close()
                return
            
            # 建立SSH连接
            self.ssh_handler = SSHHandler(
                host=host.ip,
                port=host.port,
                username=host.username,
                password=host.password,
                ssh_key=host.ssh_key
            )
            
            await asyncio.get_event_loop().run_in_executor(
                None, self.ssh_handler.connect
            )
            
            # 启动读取任务
            self.read_task = asyncio.create_task(self.read_from_ssh())
            
            await self.send(text_data=json.dumps({
                'type': 'connected',
                'message': '连接成功'
            }))
            
        except Exception as e:
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': str(e)
            }))
            await self.close()
    
    async def disconnect(self, close_code):
        """WebSocket断开"""
        if self.read_task:
            self.read_task.cancel()
        
        if self.ssh_handler:
            await asyncio.get_event_loop().run_in_executor(
                None, self.ssh_handler.close
            )
    
    async def receive(self, text_data=None, bytes_data=None):
        """接收WebSocket消息"""
        try:
            if text_data:
                data = json.loads(text_data)
                msg_type = data.get('type')
                
                if msg_type == 'stdin':
                    # 发送命令到SSH
                    command = data.get('data', '')
                    if self.ssh_handler:
                        await asyncio.get_event_loop().run_in_executor(
                            None, self.ssh_handler.send, command.encode('utf-8')
                        )
                
                elif msg_type == 'resize':
                    # 调整终端大小
                    cols = data.get('cols', 80)
                    rows = data.get('rows', 24)
                    if self.ssh_handler:
                        await asyncio.get_event_loop().run_in_executor(
                            None, self.ssh_handler.resize, cols, rows
                        )
        
        except Exception as e:
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': str(e)
            }))
    
    async def read_from_ssh(self):
        """从SSH读取数据并发送到WebSocket"""
        try:
            while self.ssh_handler and self.ssh_handler.is_active():
                data = await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_handler.recv, 1024
                )
                
                if data:
                    await self.send(text_data=json.dumps({
                        'type': 'stdout',
                        'data': data.decode('utf-8', errors='ignore')
                    }))
                else:
                    await asyncio.sleep(0.01)
        
        except asyncio.CancelledError:
            pass
        except Exception as e:
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'读取数据错误: {str(e)}'
            }))
    
    @staticmethod
    async def get_host(host_id):
        """获取主机信息（异步）"""
        from asgiref.sync import sync_to_async
        try:
            return await sync_to_async(Host.objects.get)(id=host_id)
        except Host.DoesNotExist:
            return None
